import aiomysql
import asyncio
from loguru import logger
from config import DATABASE


class MySQLHelper:
    def __init__(self):
        self.pool = None

    async def init_pool(self):
        """初始化连接池"""
        if not self.pool or self.pool.closed:
            self.pool = await aiomysql.create_pool(
                host=DATABASE['host'],
                port=DATABASE['port'],
                user=DATABASE['user'],
                password=DATABASE['password'],
                db=DATABASE['database'],
                charset="utf8mb4",
                cursorclass=aiomysql.DictCursor,
                maxsize=50,
                minsize=2,
                pool_recycle=3600,  # 连接回收时间
                echo=True,  # 打印SQL语句，方便调试
                autocommit=True,
                auth_plugin='mysql_native_password'  # 显式指定认证插件  ALTER USER 'root'@'%' IDENTIFIED WITH mysql_native_password BY '123456';FLUSH PRIVILEGES;
            )

    async def execute_query(self, query, params=None, retries=3):
        """执行查询"""
        for attempt in range(retries):
            try:
                await self.init_pool()  # 确保连接池可用
                async with self.pool.acquire() as conn:
                    async with conn.cursor() as cursor:
                        if params:
                            formatted_query = query % params
                            logger.info(f"执行查询SQL: {formatted_query}")
                        await cursor.execute(query, params)
                        return await cursor.fetchall()
            except (aiomysql.OperationalError, aiomysql.InterfaceError) as e:
                logger.error(f"查询出错，尝试重连 (尝试 {attempt + 1}): {e}")
                await asyncio.sleep(2 ** attempt)
                await self.init_pool()  # 重新初始化连接池
            except Exception as e:
                logger.error(f"查询失败: {e}")
                break
            # finally:
            #     self.pool.close()
            #     await self.pool.wait_closed()

        return None

    async def execute_update(self, query, params=None, retries=3):
        """执行更新"""
        for attempt in range(retries):
            try:
                await self.init_pool()  # 确保连接池可用
                async with self.pool.acquire() as conn:
                    async with conn.cursor() as cursor:
                        if params:
                            formatted_query = query % params
                            logger.info(f"执行修改SQL: {formatted_query}")
                        await cursor.execute(query, params)
                        await conn.commit()
                        return True
            except (aiomysql.OperationalError, aiomysql.InterfaceError) as e:
                logger.error(f"更新出错，尝试重连 (尝试 {attempt + 1}): {e}")
                await asyncio.sleep(2 ** attempt)
                await self.init_pool()  # 重新初始化连接池
            except Exception as e:
                logger.error(f"更新失败: {e}")
                break
            # finally:
            #     self.pool.close()
            #     await self.pool.wait_closed()

        return False

    async def close_pool(self):
        if self.pool and not self.pool.closed:
            self.pool.close()
            await self.pool.wait_closed()



# 数据库事务示例操作
async def example_operations():
    db = MySQLHelper()
    await db.init_pool()  # 初始化连接池

    conn = None
    try:
        # 获取连接并启动事务
        conn = await db.pool.acquire()
        async with conn.cursor() as cursor:
            await cursor.execute("START TRANSACTION;")

            # 执行第一个 SQL 操作
            await cursor.execute("UPDATE case_steps SET type = 1 WHERE id = %s;", (629,))

            # 执行第二个 SQL 操作（故意制造错误）
            await cursor.execute("UPDATE case_steps SET type = %s WHERE id = %s;", ('invalid_value', 629))

            # 提交事务
            await conn.commit()

    except Exception as e:
        logger.error(f"事务失败: {e}")

        # 如果连接存在且仍然可用，尝试回滚
        if conn and not conn.closed:
            try:
                await conn.rollback()
                logger.info("事务已回滚.")
            except Exception as rollback_error:
                logger.error(f"回滚失败: {rollback_error}")
        else:
            logger.error("无法回滚，连接已关闭或不可用.")

    finally:
        # 确保连接被释放回连接池
        if conn:
            conn.close()
            logger.info("连接已释放.")





# 使用 asyncio.run 来调用异步函数
# if __name__ == '__main__':
#     asyncio.run(example_operations())
